#!/usr/bin/env python3
"""
生成论文所需的特征重要性图表
"""

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import os
from pathlib import Path

# 设置字体和样式 - 修复方块问题
plt.rcParams['font.family'] = ['Arial', 'DejaVu Sans', 'Liberation Sans']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['figure.dpi'] = 300
sns.set_style("whitegrid")

class FeatureImportanceChartGenerator:
    def __init__(self, output_dir='Pictures'):
        """初始化图表生成器"""
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)
        
        # 解析Top 20特征数据
        self.parse_feature_data()
        
    def parse_feature_data(self):
        """解析Top 20特征重要性数据"""
        # Top 20重要特征数据 (特征名, 类型, 重要性, 有效覆盖率)
        self.top20_features = [
            ('malicious ratio', 'Graph', 0.1560),
            ('max community maliciousness ratio', 'Graph', 0.1096),
            ('sum malicious pagerank', 'Graph', 0.0993),
            ('rule PE-001', 'Rule', 0.0949),
            ('mean malicious pagerank', 'Graph', 0.0511),
            ('max malicious betweenness', 'Graph', 0.0490),
            ('mean malicious betweenness', 'Graph', 0.0488),
            ('max malicious pagerank', 'Graph', 0.0486),
            ('max degree centrality full', 'Graph', 0.0484),
            ('rule PE-002', 'Rule', 0.0411),
            ('mean degree centrality full', 'Graph', 0.0248),
            ('PE ratio', 'Graph', 0.0167),
            ('malicious node count', 'Graph', 0.0155),
            ('malicious internal edges', 'Graph', 0.0143),
            ('DT ratio', 'Graph', 0.0140),
            ('rule IG-003', 'Rule', 0.0131),
            ('max degree centrality sub', 'Graph', 0.0127),
            ('malicious density', 'Graph', 0.0123),
            ('mean degree centrality sub', 'Graph', 0.0111),
            ('category diversity', 'Graph', 0.0105)
        ]
        
        # 转换为DataFrame
        self.df_top20 = pd.DataFrame(self.top20_features, 
                                   columns=['feature_name', 'feature_type', 'importance'])

    def generate_feature_importance_chart(self):
        """生成Top 20特征重要性水平条形图 - 论文用"""
        plt.figure(figsize=(10, 8))
        
        # 准备数据
        features = self.df_top20['feature_name'].values
        importance = self.df_top20['importance'].values
        feature_types = self.df_top20['feature_type'].values
        
        # 设置颜色：图特征用蓝色，规则特征用红色
        colors = ['#1f77b4' if t == 'Graph' else '#d62728' for t in feature_types]
        
        # 创建水平条形图
        y_pos = np.arange(len(features))
        bars = plt.barh(y_pos, importance, color=colors, alpha=0.8, height=0.7)
        
        # 设置y轴标签 - 清理特征名称以提高可读性
        plt.yticks(y_pos, features, fontsize=10)
        plt.xlabel('Feature Importance', fontsize=12, fontweight='bold')
        plt.title('Top 20 Feature Importance Ranking', fontsize=14, fontweight='bold')
        
        # 反转y轴，使最重要的特征在顶部
        plt.gca().invert_yaxis()
        
        # 添加图例
        from matplotlib.patches import Patch
        legend_elements = [
            Patch(facecolor='#1f77b4', label='Graph Features'),
            Patch(facecolor='#d62728', label='Rule Features')
        ]
        plt.legend(handles=legend_elements, loc='lower right', fontsize=10)
        
        # 添加网格线
        plt.grid(axis='x', alpha=0.3, linestyle='--')
        
        # 在条形右侧添加数值标签
        for i, (bar, imp) in enumerate(zip(bars, importance)):
            plt.text(imp + 0.003, bar.get_y() + bar.get_height()/2, 
                    f'{imp:.3f}', va='center', ha='left', fontsize=8)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        output_path = self.output_dir / 'feature_importance_top20.png'
        plt.savefig(output_path, dpi=300, bbox_inches='tight', 
                   facecolor='white', edgecolor='none')
        plt.close()
        
        print(f"✅ 特征重要性图已保存: {output_path}")
        return output_path

    def generate_chart(self):
        """生成论文所需的图表"""
        print("🎨 正在生成论文特征重要性图表...")
        chart_path = self.generate_feature_importance_chart()
        
        print(f"\n📊 图表生成完成！")
        print(f"📁 保存位置: {chart_path}")
        print(f"\n💡 使用说明:")
        print(f"   - 将此图插入论文第5.2.1节，替换图\\ref{{fig:feature_importance}}的占位符")
        print(f"   - 图中蓝色表示图结构特征，红色表示规则特征")
        print(f"   - 前3名特征均为图特征，验证了PDCG结构分析的价值")

def main():
    """主函数"""
    generator = FeatureImportanceChartGenerator()
    generator.generate_chart()

if __name__ == '__main__':
    main()